"""
生成 所有品种的 期货、现货、基差、价差、期货预测、库存的一种
"""
import time
import pandas as pd
import numpy as np
from tqdm import tqdm
from copy import deepcopy
from tools.utils import load_config
from tools.utils import get_chinese_name_to_ts_code
from tools.utils import get_chinese_name_to_wind_code
from tools.utils import find_spot_or_stock_by_ts_code
from tools.utils import loc_collection, loc_zq_collection
from cal.cal_fut_spot import cal_fut_spot_diff
from cal.cal_fut import generate_main_contract_time_series
from cal.cal_fut_fut import cal_fut_fut_basis_by_ts_code
from upload.upload_future import get_main_future_contract_map
from cal.cal_pred import get_pred_by_history


def cal_six(mode: str = "期货", category: str = "tjd_category"):
    """
    :param category: 某种大宗商品分类，比如{钢铁：[螺纹钢、焦煤] ,能化:[原油、动力煤]}
    :param mode: 预测模式，期货、现货、基差、价差、期货预测、库存中的一种
    :return: 二级字典，包含期货合约中，具体分类中，具体品种的时间序列。比如:record_dic['钢铁']['螺纹钢']
    """
    pred = False
    record_dic = {}
    record_li = []

    func_map = {"期货": generate_main_contract_time_series,
                "现货": find_spot_or_stock_by_ts_code,
                "基差": cal_fut_spot_diff,
                "价差": cal_fut_fut_basis_by_ts_code,
                "库存": find_spot_or_stock_by_ts_code,
                }
    cfg_cal = load_config()
    main_contract_map = get_main_future_contract_map(cfg_cal)
    ts_code_map = get_chinese_name_to_ts_code(cfg_cal)
    wind_code_map = get_chinese_name_to_wind_code(cfg_cal)
    for category, chinese_names in tqdm(cfg_cal[category].items(), ncols=80, desc='CATEGORY'):
        name_dic = {}
        for chinese_name in chinese_names:
            ts_code = ts_code_map[chinese_name]
            func = func_map[mode]
            if mode == '期货':
                v, name_v = func(cfg_cal, ts_code, mode, main_code_li=main_contract_map.iloc[-1, :].to_list())
            else:
                v, name_v = func(cfg_cal, ts_code, mode)
            name_dic[chinese_name] = v
            save_dic = {
                'tjd_category': category,
                'chinese_name': chinese_name,
                'mode': mode,
                'ts_code': ts_code,
                'symbol': ts_code.split('_')[0].upper(),
                'wind_code': wind_code_map[chinese_name],
                'pred': pred,
                'default': False,
                'meiosis': '-',
                'minuend': '-',
                'seasonal': False,
            }
            if len(v) > 0:
                many_li = write_dic_to_excel(v, name_v, deepcopy(save_dic))
                record_li.extend(many_li)
                # 增加周和月频率
                for dic in many_li:
                    dic['freq'] = detect_freq(dic)
                    record_li.extend(day_to_week_or_month(dic, '周'))
                    record_li.extend(day_to_week_or_month(dic, '月'))
            else:
                print(f"miss...{chinese_name}_{mode}")
        record_dic[category] = name_dic
    return record_dic, record_li


def write_dic_to_excel(dic: dict, name_dic: dict, base_dic: dict):
    li = []
    set_default = False
    for i, (k, v) in enumerate(dic.items()):
        base_dic_cp = deepcopy(base_dic)
        v = pd.DataFrame(v)
        name = v.columns.to_list()[0]
        base_dic_cp['columns'] = v.columns.to_list()
        base_dic_cp['t'] = v.index.to_list()
        base_dic_cp['value'] = v.values.tolist()
        if base_dic_cp['mode'] in ['基差', '价差']:
            base_dic_cp['meiosis'] = name_dic[k]['meiosis']
            base_dic_cp['minuend'] = name_dic[k]['minuend']
        if base_dic_cp['mode'] in ['基差', '期货']:

            if ('主力' in name) and not set_default:
                base_dic_cp['default'] = True
                set_default = True
        elif base_dic_cp['mode'] in ['价差']:
            if ('主力' in name) and not set_default:
                base_dic_cp['default'] = True
                set_default = True
        else:
            if i < 1:
                base_dic_cp['default'] = True
        li.append(base_dic_cp)
    return li


def to_seasonal(df: pd.DataFrame, mode="现货", freq="周"):
    freq_map = {'日': "day", '周': 'week', '月': 'month'}
    freq_ = freq_map[freq]
    df = pd.DataFrame(df)
    name = df.columns.to_list()[0]
    df.index = [pd.to_datetime(str(s)) for s in df.index]
    df['year'] = df.index.year
    if freq == "周":
        df['week'] = df.index.isocalendar().week
    elif freq == '月':
        df['month'] = df.index.month
    else:
        df['day'] = df.index.dayofyear
    li = []
    for year, year_df in df.groupby(['year']):
        w_dic = {}
        for d_w_m, week_df in year_df.groupby([freq_]):
            w_dic[d_w_m] = week_df.mean(axis=0)[name]
        se = pd.Series(w_dic, name=year)
        li.append(se)
    data = pd.concat(li, axis=1)
    data.sort_index(inplace=True)
    return data


def detect_freq(dic: dict):
    t_li = dic['t']
    if len(t_li) in [53, 54]:
        dic['freq'] = '周'
    else:
        ts = [pd.to_datetime(s) for s in t_li]
        length_t = ts[-1] - ts[0]
        day_cal = length_t.days
        day_real = len(t_li)
        ratio = day_real / day_cal
        if ratio > 0.56:
            dic['freq'] = "日"
        elif ratio > 0.13:
            dic['freq'] = '周'
        # elif ratio > 0.09:
        #     dic['freq'] = '旬'
        else:
            dic['freq'] = '月'
    return dic['freq']


def day_to_week_or_month(dic, mode='周'):
    if dic['freq'] not in ['周', '月']:
        t = dic['t']
        column = dic['columns']
        data = pd.DataFrame(dic['value'], index=t, columns=column)
        data.index = pd.to_datetime(data.index)
        data['year'] = data.index.isocalendar().year
        if mode == '周':
            data['周'] = data.index.isocalendar().week
        else:
            data['月'] = data.index.month
        data['key'] = [f"{a}_{b}" for a, b in zip(data[mode].to_list(), data['year'].to_list())]
        li = []
        for week, mode_df in data.groupby(['key']):
            week_df = mode_df.dropna(axis=0, how='all')
            if len(week_df) > 0:
                fit_df = week_df
            else:
                fit_df = mode_df
            name = fit_df.index.to_list()[-1]
            v_se = fit_df.iloc[-1, :]
            v_se.name = name
            li.append(v_se)
        transform_data = pd.concat(li, axis=1).T
        transform_data.sort_index(inplace=True)
        transform_data.index = [str(s)[:10] for s in transform_data.index]
        new_data = transform_data.loc[:, column]
        new_dic = deepcopy(dic)
        new_dic['t'] = new_data.index.to_list()
        new_dic['columns'] = column
        new_dic['value'] = new_data.values.tolist()
        new_dic['freq'] = mode
        return [new_dic]
    else:
        return []


def find_history_and_upload_pred(coll_select, cfg, freq='日'):
    c_mp = get_main_future_contract_map(cfg)
    c_li = [arr.unique().tolist()[-1] for k, arr in c_mp.iteritems()]
    data = list(coll_select.find({'mode': '期货'}))
    new_data = [v for v in data if (v['columns'][0].replace('_主力', '') in c_li and v['freq'] == freq)]
    pred_data = []
    log_li = []
    for v in new_data:
        result_v = get_pred_by_history(v)
        if len(result_v) > 0:
            dic = result_v[0]
            chinese_name = dic['chinese_name']
            col = dic['columns'][0]
            log_li.append(f"名称:{chinese_name},主力合约:{col},{freq}度预测上传成功")
        pred_data.extend(result_v)
    coll_select.insert_many(pred_data)
    return pred_data, log_li


if __name__ == "__main__":
    t1 = time.time()
    cfg = load_config()
    coll = loc_collection('to_it', '2022-01-12')
    for mode_ in [
        '现货',
        '期货',
        '基差',
        '价差',
        '库存'
    ]:
        record_d, record_l = cal_six(mode_, "tjd_category")
        coll.insert_many(record_l)
    t2 = time.time()
    t = t2 - t1
    find_history_and_upload_pred(coll, cfg,'日')
    # find_history_and_upload_pred(coll, '周')
